# deepseek_mcp_agent.py
import ollama
import json
import re
from mcp_tools import MCPWeatherTool, MCPCalculatorTool, MCPFileTool
from mcp_tools import MCPWebSearchTool

class DeepSeekMCPAgent:
    def __init__(self, model="deepseek-r1:14b"):
        self.model = model
        self.memory = []
        self.tools = {
            'weather': MCPWeatherTool(),
            'calculator': MCPCalculatorTool(),
            'search': MCPWebSearchTool(),
            'file': MCPFileTool()
        }
        
        # MCP工具描述，用于提示词
        self.tool_descriptions = self._build_tool_descriptions()
    
    def _build_tool_descriptions(self):
        """构建工具描述供模型使用"""
        return """
可用工具列表：
1. weather - 查询天气信息
   使用格式: <tool name="weather">{"location": "城市名称"}</tool>
   
2. calculator - 执行数学计算
   使用格式: <tool name="calculator">{"expression": "数学表达式"}</tool>
   
3. search - 网络搜索
   使用格式: <tool name="search">{"query": "搜索关键词"}</tool>
   
4. file - 文件操作
   读取文件: <tool name="file">{"action": "read", "filepath": "文件路径"}</tool>
   写入文件: <tool name="file">{"action": "write", "filepath": "文件路径", "content": "内容"}</tool>
"""
    
    def _call_llm(self, prompt):
        """调用Ollama模型"""
        messages = self.memory.copy()
        messages.append({"role": "user", "content": prompt})
        
        response = ollama.chat(
            model=self.model,
            messages=messages,
            options={'temperature': 0.3}
        )
        return response['message']['content']
    
    def _extract_tool_call(self, text):
        """从模型响应中提取工具调用信息"""
        tool_pattern = r'<tool name="(\w+)">(.*?)</tool>'
        matches = re.findall(tool_pattern, text, re.DOTALL)
        
        if matches:
            tool_name, params_str = matches[0]
            try:
                params = json.loads(params_str)
                return tool_name, params
            except json.JSONDecodeError:
                return None, None
        return None, None
    
    def _execute_tool(self, tool_name, params):
        """执行MCP工具"""
        if tool_name in self.tools:
            tool = self.tools[tool_name]
            
            if tool_name == 'weather':
                return tool.get_weather(params.get('location', ''))
            elif tool_name == 'calculator':
                return tool.calculate(params.get('expression', ''))
            elif tool_name == 'search':
                return tool.search(params.get('query', ''))
            elif tool_name == 'file':
                action = params.get('action', '')
                if action == 'read':
                    return tool.read_file(params.get('filepath', ''))
                elif action == 'write':
                    return tool.write_file(
                        params.get('filepath', ''), 
                        params.get('content', '')
                    )
        return {'error': f'未知工具: {tool_name}', 'success': False}
    
    def process_query(self, user_input):
        """处理用户查询"""
        # 构建包含工具描述的提示词
        prompt = f"""
{self.tool_descriptions}

用户问题: {user_input}

请分析用户需求，如果需要使用工具，请按照指定格式调用工具。
如果不需要工具，请直接回答用户问题。

思考过程:
1. 分析用户意图
2. 决定是否需要工具
3. 如果需要，选择合适的工具并生成正确的参数格式
"""
        # 调用模型获取初始响应
        initial_response = self._call_llm(prompt)
        print(f"模型初始响应: {initial_response}")
        
        # 检查是否需要工具调用
        tool_name, params = self._extract_tool_call(initial_response)
        
        if tool_name and params:
            print(f"检测到工具调用: {tool_name}, 参数: {params}")
            
            # 执行工具
            tool_result = self._execute_tool(tool_name, params)
            print(f"工具执行结果: {tool_result}")
            
            # 将工具结果反馈给模型进行最终回答
            followup_prompt = f"""
用户原始问题: {user_input}
工具调用结果: {json.dumps(tool_result, ensure_ascii=False)}

请根据工具执行结果生成最终回答。
"""
            final_response = self._call_llm(followup_prompt)
            
            # 更新记忆
            self.memory.extend([
                {"role": "user", "content": user_input},
                {"role": "assistant", "content": final_response}
            ])
            
            return {
                'final_response': final_response,
                'tool_used': tool_name,
                'tool_result': tool_result,
                'success': True
            }
        else:
            # 不需要工具，直接返回回答
            self.memory.extend([
                {"role": "user", "content": user_input},
                {"role": "assistant", "content": initial_response}
            ])
            
            return {
                'final_response': initial_response,
                'tool_used': None,
                'success': True
            }
    
    def run_interactive(self):
        """运行交互式对话"""
        print("DeepSeek MCP Agent 已启动! (输入 'exit' 退出)")
        print("=" * 50)
        
        while True:
            try:
                user_input = input("\n您: ").strip()
                
                if user_input.lower() in ['exit', 'quit', '退出']:
                    print("再见!")
                    break
                
                if not user_input:
                    continue
                
                print("Agent思考中...")
                result = self.process_query(user_input)
                
                print(f"\nDeepSeek: {result['final_response']}")
                if result['tool_used']:
                    print(f"[使用的工具: {result['tool_used']}]")
                    
            except KeyboardInterrupt:
                print("\n\n会话结束!")
                break
            except Exception as e:
                print(f"发生错误: {e}")

# 使用示例
if __name__ == "__main__":
    # 创建Agent实例
    agent = DeepSeekMCPAgent()
    
    # 测试示例
    test_queries = [
        "北京今天的天气怎么样？",
        "计算一下 125 * 8 + 300 等于多少？",
        "搜索一下人工智能的最新发展",
        "创建一个文件并写入'Hello MCP World!'"
    ]
    
    print("测试MCP Agent功能:")
    for query in test_queries:
        print(f"\n测试查询: {query}")
        result = agent.process_query(query)
        print(f"响应: {result['final_response']}")
        if result['tool_used']:
            print(f"使用的工具: {result['tool_used']}")
    
    # 启动交互式对话
    agent.run_interactive()